import numpy as np
import acl
import hccl

# 初始化 ACL
acl.init()

# 设备数量
device_num = 2

# 初始化 HCCL
hccl.hccl_init(0)

# 创建通信组
world_size = device_num
rank = 0  # 当前进程的 rank
comm = hccl.create_group(world_size, rank)

# 初始化数据
data = np.ones((10,), dtype=np.float32)
data_buf = acl.rt.malloc(data.nbytes, acl.rt.MemMallocPolicy.MALLOC_HUGE_FIRST)
acl.rt.memcpy(data_buf, data.nbytes, data.ctypes.data, data.nbytes, acl.rt.MemcpyKind.MEMCPY_HOST_TO_DEVICE)

# 执行 AllReduce 操作
hccl.all_reduce(data_buf, data_buf, 10, hccl.DataType.DT_FLOAT, hccl.ReduceOp.REDUCE_SUM, comm)

# 同步设备
acl.rt.synchronize_device()

# 将结果拷贝回主机
result = np.empty((10,), dtype=np.float32)
acl.rt.memcpy(result.ctypes.data, result.nbytes, data_buf, result.nbytes, acl.rt.MemcpyKind.MEMCPY_DEVICE_TO_HOST)

print("AllReduce 结果:", result)

# 释放资源
acl.rt.free(data_buf)
hccl.destroy_group(comm)
acl.finalize()